# coding=utf-8
# Chern-C
# Copyright (C) 2024 NJU-SAML
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import asyncio
import logging
import os
import shutil
import tempfile
from pathlib import Path
from typing import List, Optional, Set, Union

from chernc.agents.translator import TranslatorBase
from chernc.autobuild.clang_build import create_build_script
from chernc.codeql.codeql_database import CodeqlDatabase
from chernc.codeql.codeql_engine import CodeqlEngine
from chernc.compiler.rust.error_messages import CargoMessageCompilerMessage, RustcErrorSpan
from chernc.config import CherncConfig
from chernc.constants import DEFAULT_LOGGER_DIR, DEFAULT_LOGGER_FILE_NAME, DEFAULT_PROJECT_REPAIR_DIR
from chernc.core.cargo import CargoConfig, Package
from chernc.core.cargo_manager import CargoManager
from chernc.core.interact import InteractEngine
from chernc.core.project_repair import ProjectRepairEngine
from chernc.core.repair import filter_error_messages
from chernc.core.result import DummySnippetResult, FileResult, ProjectRepairResult, ProjectResult, SnippetResult
from chernc.logging import build_logger
from chernc.slicer.code_slice import CodeSlice
from chernc.utils.hash_utils import calculate_md5

# from chernc.autobuild.detector import BuildSystemDetector


logger: logging.Logger = build_logger(__name__, logger_filename=DEFAULT_LOGGER_FILE_NAME, logger_dir=DEFAULT_LOGGER_DIR)


class ProjectTranspiler(object):
    """
    项目翻译器，用于存储项目级别的状态。
    """

    def __init__(
        self,
        config: CherncConfig,
        project_path: Union[str, Path],
        database_path: Union[str, Path],
        translator: TranslatorBase,
        codeql_bin: str = "codeql",
    ) -> None:
        if isinstance(project_path, Path):
            project_path = str(project_path)
        if isinstance(database_path, Path):
            database_path = str(database_path)
        # 配置信息
        self.config: CherncConfig = config
        # 项目基本信息
        self.project_path: str = os.path.abspath(project_path)  # 请注意，这里我们要保证project_path是绝对路径，以确保唯一性
        self.project_hash: str = calculate_md5(os.path.abspath(project_path))
        self.project_name: str = Path(project_path).name

        # 静态分析数据库路径
        self.database_path: str = database_path
        # 静态分析引擎
        self.codeql_engine: CodeqlEngine = CodeqlEngine(codeql_bin=codeql_bin)
        # 翻译API
        self.translator: TranslatorBase = translator
        # self.build_system_detector = BuildSystemDetector()

    def _redirect_path_to_output(self, source_path: str, output_path: str) -> str:
        # check if source_path is under self.project_path
        if not os.path.abspath(source_path).startswith(os.path.abspath(self.project_path)):
            raise ValueError("source_path is not under the project path")
        rel_filepath = os.path.relpath(source_path, self.project_path)
        rel_filedir = os.path.dirname(rel_filepath)
        return os.path.join(output_path, rel_filedir, rel_filepath)

    def _replace_suffix_c2rust(self, filepath: str) -> str:
        return Path(filepath).with_suffix(".rs").as_posix()

    def write_project(
        self, output_path: str, project_result: Union[ProjectResult, ProjectRepairResult], with_comments: bool = False
    ) -> None:
        # clear the source path
        # shutil.rmtree(os.path.join(output_path, self.config.source_path), ignore_errors=True)
        # Create Cargo.toml
        cargo_manager = CargoManager.from_cargo_config(project_result.cargo, output_path)
        # Flush Now to avoid the bin section is missing in the cargo.toml when we check it
        cargo_manager.flush()

        # Write Files
        for file_path, file_result in project_result.files.items():
            rust_file_name = os.path.basename(file_path)
            rel_filedir = os.path.dirname(file_path)

            # out_file = os.path.join(output_path, self.config.source_path, rel_filedir, rust_file_name)
            # out_filedir = os.path.join(output_path, self.config.source_path, rel_filedir)
            out_file = os.path.join(output_path, rel_filedir, rust_file_name)
            out_filedir = os.path.join(output_path, rel_filedir)
            if not os.path.exists(out_filedir):
                os.makedirs(out_filedir, exist_ok=True)
            if with_comments:
                text = file_result.get_text_with_details(
                    mark_position=self.config.mark_translation_position, mark_source=self.config.mark_translation_source
                )
            else:
                text = file_result.text
            with open(out_file, "w", encoding="utf-8") as f:
                f.write(text)

        # Write Cargo.toml
        cargo_manager.close()

    def _analyze_source(self):
        has_build_system: bool = False
        build_script_path: Optional[str] = None
        if not has_build_system:
            build_script_path = os.path.abspath(os.path.join(self.project_path, "chernc_build.sh"))
            build_script = create_build_script(self.project_path)
            with open(build_script_path, "w", encoding="utf-8") as f:
                f.write(build_script)

        else:
            # TODO: Add build script locator
            build_script_path = None

        # 静态分析数据库
        self.codeql_database = CodeqlDatabase(
            codeql_engine=self.codeql_engine,
            project_path=self.project_path,
            database_path=self.database_path,
            build_script_path=build_script_path,
        )

    def translate_project(self, output_path: str) -> None:
        self._analyze_source()
        if not os.path.exists(output_path):
            os.makedirs(output_path, exist_ok=True)

        # 项目级别翻译
        os.makedirs(os.path.join(output_path, "chernc_original"), exist_ok=True)
        with tempfile.TemporaryDirectory() as tmp_dir, CargoManager(
            project_path=tmp_dir,
            dependencies={"libc": "0.2.162"},
        ) as package_installer:
            os.makedirs(os.path.join(tmp_dir, "src"), exist_ok=True)
            with open(os.path.join(tmp_dir, "src", "main.rs"), "w") as f:
                f.write('fn main() {\n    println!("Hello, world!");\n}\n')
            engine = InteractEngine(self, self.translator, package_installer)
            project_result = asyncio.run(engine.async_trans_project(self, out_path=output_path))

        # 记录最初版
        os.makedirs(os.path.join(output_path, "chernc_original"), exist_ok=True)
        self.write_project(
            output_path=os.path.join(output_path, "chernc_original"), project_result=project_result, with_comments=True
        )

        # 项目级别修复
        project_repair_path = os.path.join(output_path, DEFAULT_PROJECT_REPAIR_DIR)
        repair_engine = ProjectRepairEngine(self, self.translator)
        if os.path.exists(project_repair_path):
            shutil.rmtree(project_repair_path)
        repaired_project_result = asyncio.run(
            repair_engine.repair_project(
                project_result, temp_dir=project_repair_path, max_fix_count=self.config.max_project_repair
            )
        )

        # 写入最终文件
        self.write_project(output_path=output_path, project_result=repaired_project_result, with_comments=True)

        self.count_compiled_pass(repaired_project_result)

    def count_compiled_pass(self, project_result: ProjectRepairResult) -> None:
        """
        Count the number of slices that can be compiled successfully.
        """
        all_snippets = [snippet for file in project_result.files.values() for snippet in file.snippets]
        with tempfile.TemporaryDirectory() as tmp_dir, CargoManager.from_cargo_config(
            project_result.cargo, tmp_dir
        ) as cargo_manager:
            self.write_project(tmp_dir, project_result, with_comments=False)
            messages = [m.message for m in cargo_manager.build() if isinstance(m, CargoMessageCompilerMessage)]
        messages = filter_error_messages(messages, anonymous=False)
        spans: List[RustcErrorSpan] = []
        for m in messages:
            for span in m.spans:
                while span.expansion:
                    span = span.expansion.span
                spans.append(span)
        err_snippets: Set[SnippetResult] = set()
        for span in spans:
            norm_path = os.path.normpath(span.file_name).replace("\\", "/")
            # file_path = os.path.relpath(norm_path, self.config.source_path)
            file_path = norm_path
            file_result = project_result.files[file_path]
            found_snippet = file_result.find_snippet_by_line(span.line_start - 1)
            if found_snippet:
                err_snippets.add(found_snippet)

        # collect all the slices
        total = sum(len(snippet.slices) for snippet in all_snippets)
        err = sum(len(snippet.slices) for snippet in err_snippets)
        dummy = sum(len(snippet.slices) for snippet in all_snippets if isinstance(snippet, DummySnippetResult))
        logger.info(f"Compilation pass count: {total - err}/{total}")

        logger.info(f"Transpiled slices count: {total - dummy}/{total}")
        # collect slices in src
        total = sum(
            1 for snippet in all_snippets for slice in snippet.slices if not os.path.basename(slice.file).startswith("test-")
        )
        err = sum(
            1 for snippet in err_snippets for slice in snippet.slices if not os.path.basename(slice.file).startswith("test-")
        )
        logger.info(f"Compilation pass in source files count: {total - err}/{total}")

        # collect slices in test
        total = sum(1 for snippet in all_snippets for slice in snippet.slices if os.path.basename(slice.file).startswith("test-"))
        err = sum(1 for snippet in err_snippets for slice in snippet.slices if os.path.basename(slice.file).startswith("test-"))
        logger.info(f"Compilation pass in test files count: {total - err}/{total}")

        def is_function(slice: CodeSlice) -> bool:
            return any(decl[0].endswith("Function") for decl in slice.decls)

        # collect slices of functions
        total = sum(1 for snippet in all_snippets for slice in snippet.slices if is_function(slice))
        err = sum(1 for snippet in err_snippets for slice in snippet.slices if is_function(slice))
        logger.info(f"Compilation pass of functions count: {total - err}/{total}")
        return


if __name__ == "__main__":
    config = CherncConfig()
    translator = TranslatorBase()
    transpiler = ProjectTranspiler(config, "", "", translator)
    slice = CodeSlice(None, "main.c", 1, 1)
    slice.text = "int main() { return 0; }"
    slice2 = CodeSlice(None, "main.c", 2, 2)
    slice2.text = "int m() { return 0; }"
    snippet = SnippetResult([slice], 'fn main() {\n    println!("Hello, world!");\n}\n', ["main"])
    snippet2 = SnippetResult([slice2], "fn m", [])
    file_result = FileResult("main.rs")
    file_result.add_snippet(snippet)
    file_result.add_snippet(snippet2)
    cargo_config = CargoConfig(package=Package(name="test", version="0.1.0", authors=[]))
    project_result = ProjectResult(cargo_config, {"main.rs": file_result})
    transpiler.count_compiled_pass(project_result.to_project_repair_result())
